import os
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import cv2
from CNNImageClassifyTrain import CNN, ImageDataset

model = CNN()
model.load_state_dict(torch.load('cnn_model.pth'))
model.eval()

test_data_dir = './test'
test_dataset = ImageDataset(test_data_dir, transform=None)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

if __name__ == '__main__':
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            print(predicted)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print(f'Accuracy of the network on the test images: {100 * correct / total}%')
